/**
 * React Hook for streaming AI content generation.
 * Supports both WebSocket and SSE with automatic fallback.
 */

import { useState, useEffect, useRef, useCallback } from 'react';
import { message } from 'antd';
import { GenerationWebSocketService, createGenerationWebSocket } from '../services/generation.websocket.service';
import { GenerationSSEService, createGenerationSSE } from '../services/sse.service';

export interface UseStreamGenerationOptions {
  /** Use SSE instead of WebSocket */
  useSSE?: boolean;
  /** Authentication token */
  token?: string;
  /** Callback when generation starts */
  onStart?: (metadata: any) => void;
  /** Callback for each token received */
  onToken?: (token: string) => void;
  /** Callback when generation completes */
  onComplete?: (metadata: any) => void;
  /** Callback on error */
  onError?: (error: string) => void;
  /** Auto-reconnect on disconnect */
  autoReconnect?: boolean;
}

export interface GenerationParams {
  templateId: string;
  variables: Record<string, any>;
  model?: string;
}

export interface UseStreamGenerationResult {
  /** Current generated content */
  content: string;
  /** Whether generation is in progress */
  isGenerating: boolean;
  /** Whether connected to streaming service */
  isConnected: boolean;
  /** Error message if any */
  error: string | null;
  /** Generation metadata */
  metadata: any | null;
  /** Start generation */
  startGeneration: (params: GenerationParams) => Promise<void>;
  /** Cancel ongoing generation */
  cancelGeneration: () => void;
  /** Clear content */
  clearContent: () => void;
  /** Manually connect to streaming service */
  connect: () => Promise<void>;
  /** Manually disconnect from streaming service */
  disconnect: () => void;
}

/**
 * Hook for streaming AI content generation.
 */
export function useStreamGeneration(
  taskId: string,
  options: UseStreamGenerationOptions = {}
): UseStreamGenerationResult {
  const {
    useSSE = false,
    token,
    onStart,
    onToken,
    onComplete,
    onError,
    autoReconnect = true,
  } = options;

  // State
  const [content, setContent] = useState<string>('');
  const [isGenerating, setIsGenerating] = useState(false);
  const [isConnected, setIsConnected] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [metadata, setMetadata] = useState<any | null>(null);

  // Refs
  const serviceRef = useRef<GenerationWebSocketService | GenerationSSEService | null>(null);
  const contentBufferRef = useRef<string>('');

  /**
   * Initialize streaming service.
   */
  const initializeService = useCallback(() => {
    if (serviceRef.current) {
      return;
    }

    if (useSSE) {
      // SSE will be created when starting generation (needs templateId)
      return;
    } else {
      // Create WebSocket service
      const wsService = createGenerationWebSocket(taskId, token);
      
      // Setup connection handlers
      wsService.on('connected', () => {
        setIsConnected(true);
        setError(null);
        console.log('WebSocket connected for generation');
      });

      wsService.on('disconnected', () => {
        setIsConnected(false);
        if (isGenerating) {
          setError('Connection lost during generation');
        }
      });

      wsService.on('error', (err: any) => {
        console.error('WebSocket error:', err);
        setError('Connection error occurred');
        onError?.('Connection error occurred');
      });

      // Setup generation handlers
      wsService.setGenerationHandlers({
        onStart: (meta) => {
          setIsGenerating(true);
          setError(null);
          setMetadata(meta);
          contentBufferRef.current = '';
          onStart?.(meta);
        },
        onToken: (token) => {
          contentBufferRef.current += token;
          setContent(contentBufferRef.current);
          onToken?.(token);
        },
        onComplete: (meta) => {
          setIsGenerating(false);
          setMetadata(meta);
          onComplete?.(meta);
        },
        onError: (err) => {
          setIsGenerating(false);
          setError(err);
          onError?.(err);
          message.error(`Generation error: ${err}`);
        },
      });

      serviceRef.current = wsService;
    }
  }, [taskId, token, useSSE, onStart, onToken, onComplete, onError]);

  /**
   * Connect to streaming service.
   */
  const connect = useCallback(async () => {
    try {
      if (!serviceRef.current) {
        initializeService();
      }

      if (serviceRef.current && !useSSE) {
        await serviceRef.current.connect();
      }
    } catch (err) {
      console.error('Failed to connect:', err);
      setError('Failed to connect to streaming service');
      
      // Try SSE fallback if WebSocket fails
      if (!useSSE && autoReconnect) {
        message.warning('WebSocket connection failed, trying SSE fallback...');
        // Note: In production, you might want to implement automatic fallback
      }
    }
  }, [initializeService, useSSE, autoReconnect]);

  /**
   * Disconnect from streaming service.
   */
  const disconnect = useCallback(() => {
    if (serviceRef.current) {
      serviceRef.current.disconnect();
      serviceRef.current = null;
      setIsConnected(false);
    }
  }, []);

  /**
   * Start content generation.
   */
  const startGeneration = useCallback(async (params: GenerationParams) => {
    try {
      setError(null);
      contentBufferRef.current = '';
      setContent('');

      if (useSSE) {
        // Create SSE service for this generation
        const sseService = createGenerationSSE(taskId, params.templateId, params.model);
        
        // Setup handlers
        sseService.setGenerationHandlers({
          onStart: (meta) => {
            setIsGenerating(true);
            setIsConnected(true);
            setMetadata(meta);
            onStart?.(meta);
          },
          onToken: (token) => {
            contentBufferRef.current += token;
            setContent(contentBufferRef.current);
            onToken?.(token);
          },
          onComplete: (meta) => {
            setIsGenerating(false);
            setMetadata(meta);
            onComplete?.(meta);
            // Disconnect SSE after completion
            sseService.disconnect();
            setIsConnected(false);
          },
          onError: (err) => {
            setIsGenerating(false);
            setError(err);
            onError?.(err);
            message.error(`Generation error: ${err}`);
            sseService.disconnect();
            setIsConnected(false);
          },
        });

        serviceRef.current = sseService;
        await sseService.start();
      } else {
        // WebSocket generation
        if (!serviceRef.current || !isConnected) {
          await connect();
        }

        if (serviceRef.current && 'startGeneration' in serviceRef.current) {
          await serviceRef.current.startGeneration({
            template_id: params.templateId,
            variables: params.variables,
            model: params.model,
          });
        }
      }
    } catch (err) {
      console.error('Failed to start generation:', err);
      setError('Failed to start generation');
      setIsGenerating(false);
      message.error('Failed to start generation');
    }
  }, [taskId, useSSE, isConnected, connect, onStart, onToken, onComplete, onError]);

  /**
   * Cancel ongoing generation.
   */
  const cancelGeneration = useCallback(() => {
    if (serviceRef.current && isGenerating) {
      if ('cancelGeneration' in serviceRef.current) {
        serviceRef.current.cancelGeneration();
      }
      setIsGenerating(false);
      message.info('Generation cancelled');
    }
  }, [isGenerating]);

  /**
   * Clear generated content.
   */
  const clearContent = useCallback(() => {
    setContent('');
    contentBufferRef.current = '';
    setMetadata(null);
    setError(null);
  }, []);

  // Initialize service on mount
  useEffect(() => {
    if (!useSSE) {
      initializeService();
    }

    return () => {
      disconnect();
    };
  }, []);

  // Auto-connect for WebSocket
  useEffect(() => {
    if (!useSSE && autoReconnect && !isConnected && !isGenerating) {
      connect();
    }
  }, [useSSE, autoReconnect, isConnected, isGenerating, connect]);

  return {
    content,
    isGenerating,
    isConnected,
    error,
    metadata,
    startGeneration,
    cancelGeneration,
    clearContent,
    connect,
    disconnect,
  };
}

export default useStreamGeneration;